{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from json import load\n",
    "import os\n",
    "import argparse\n",
    "import random\n",
    "from copy import deepcopy\n",
    "from munch import Munch\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from fedlab.utils.aggregator import Aggregators\n",
    "from fedlab.utils.serialization import SerializationTool\n",
    "from fedlab.utils.functional import evaluate, get_best_gpu\n",
    "\n",
    "from fedlab.models.mlp import MLP\n",
    "from fedlab.models.cnn import CNN_MNIST\n",
    "from fedlab.contrib.algorithm.basic_server import SyncServerHandler\n",
    "from fedlab.contrib.algorithm.basic_client import SGDSerialClientTrainer\n",
    "from fedlab.contrib.dataset.pathological_mnist import PathologicalMNIST\n",
    "from fedlab.contrib.dataset.partitioned_mnist import PartitionedMNIST\n",
    "\n",
    "from fedlab.utils.functional import evaluate, setup_seed\n",
    "from fedlab.contrib.algorithm.fedprox import FedProxServerHandler, FedProxSerialClientTrainer\n",
    "from fedlab.contrib.algorithm.scaffold import ScaffoldSerialClientTrainer, ScaffoldServerHandler\n",
    "from fedlab.contrib.algorithm.fednova import FedNovaSerialClientTrainer, FedNovaServerHandler\n",
    "from fedlab.contrib.algorithm.feddyn import FedDynSerialClientTrainer, FedDynServerHandler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Munch()\n",
    "args.total_client = 100\n",
    "args.com_round = 100\n",
    "args.sample_ratio = 0.5\n",
    "args.batch_size = 100\n",
    "args.epochs = 5\n",
    "args.lr = 0.05\n",
    "\n",
    "args.preprocess = False\n",
    "args.seed = 0\n",
    "\n",
    "args.alg = \"fedavg\"  # fedavg, fedprox, scaffold, fednova, feddyn\n",
    "# optim parameter\n",
    "\n",
    "args.mu = 0.1 # fedprox\n",
    "args.alpha = 0.01 # feddyn\n",
    "\n",
    "setup_seed(args.seed)\n",
    "test_data = torchvision.datasets.MNIST(root=\"../datasets/mnist/\",\n",
    "                                       train=False,\n",
    "                                       transform=transforms.ToTensor())\n",
    "\n",
    "test_loader = DataLoader(test_data, batch_size=1024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MLP(784, 10)\n",
    "# model = CNN_MNIST()\n",
    "\n",
    "if args.alg == \"fedavg\":\n",
    "    handler = SyncServerHandler(model=model, global_round=args.com_round, sample_ratio=args.sample_ratio)\n",
    "    trainer = SGDSerialClientTrainer(model, args.total_client, cuda=True)\n",
    "    trainer.setup_optim(args.epochs, args.batch_size, args.lr)\n",
    "\n",
    "if args.alg == \"fedprox\":\n",
    "    handler = FedProxServerHandler(model=model, global_round=args.com_round, sample_ratio=args.sample_ratio)\n",
    "    trainer = FedProxSerialClientTrainer(model, args.total_client, cuda=True)\n",
    "    trainer.setup_optim(args.epochs, args.batch_size, args.lr, mu=args.mu)\n",
    "\n",
    "if args.alg == \"scaffold\":\n",
    "    handler = ScaffoldServerHandler(model=model, global_round=args.com_round, sample_ratio=args.sample_ratio)\n",
    "    handler.setup_optim(lr=args.lr)\n",
    "\n",
    "    trainer = ScaffoldSerialClientTrainer(model, args.total_client, cuda=True)\n",
    "    trainer.setup_optim(args.epochs, args.batch_size, args.lr)\n",
    "\n",
    "if args.alg == \"fednova\":\n",
    "    handler = FedNovaServerHandler(model=model, global_round=args.com_round, sample_ratio=args.sample_ratio)\n",
    "    handler.setup_optim()\n",
    "    trainer = FedNovaSerialClientTrainer(model, args.total_client, cuda=True)\n",
    "    trainer.setup_optim(args.epochs, args.batch_size, args.lr)\n",
    "\n",
    "if args.alg == \"feddyn\":\n",
    "    handler = FedDynServerHandler(model=model, global_round=args.com_round, sample_ratio=args.sample_ratio)\n",
    "    handler.setup_optim(alpha=args.alpha)\n",
    "    trainer = FedDynSerialClientTrainer(model, args.total_client, cuda=True)\n",
    "    trainer.setup_optim(args.epochs, args.batch_size, args.lr, args.alpha)\n",
    "\n",
    "# mnist = PathologicalMNIST(root='./datasets/mnist/', path=\"./datasets/mnist/pathmnist\", num_clients=args.total_client, shards=200)\n",
    "mnist = PartitionedMNIST(root='../datasets/mnist/',\n",
    "                         path=\"../datasets/mnist/fedmnist_iid\",\n",
    "                         num_clients=args.total_client,\n",
    "                         partition=\"iid\", \n",
    "                         dir_alpha=args.alpha,\n",
    "                         preprocess=args.preprocess,\n",
    "                         transform=transforms.Compose(\n",
    "                             [transforms.ToPILImage(), transforms.ToTensor()]))\n",
    "# mnist.preprocess()\n",
    "trainer.setup_dataset(mnist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Round 1, Test Accuracy: 0.4346, Max Acc: 0.4346\n",
      "Round 2, Test Accuracy: 0.6325, Max Acc: 0.6325\n",
      "Round 3, Test Accuracy: 0.7072, Max Acc: 0.7072\n",
      "Round 4, Test Accuracy: 0.7555, Max Acc: 0.7555\n",
      "Round 5, Test Accuracy: 0.7980, Max Acc: 0.7980\n",
      "Round 6, Test Accuracy: 0.8210, Max Acc: 0.8210\n",
      "Round 7, Test Accuracy: 0.8428, Max Acc: 0.8428\n",
      "Round 8, Test Accuracy: 0.8578, Max Acc: 0.8578\n",
      "Round 9, Test Accuracy: 0.8697, Max Acc: 0.8697\n",
      "Round 10, Test Accuracy: 0.8764, Max Acc: 0.8764\n",
      "Round 11, Test Accuracy: 0.8844, Max Acc: 0.8844\n",
      "Round 12, Test Accuracy: 0.8881, Max Acc: 0.8881\n",
      "Round 13, Test Accuracy: 0.8918, Max Acc: 0.8918\n",
      "Round 14, Test Accuracy: 0.8948, Max Acc: 0.8948\n",
      "Round 15, Test Accuracy: 0.8974, Max Acc: 0.8974\n",
      "Round 16, Test Accuracy: 0.8994, Max Acc: 0.8994\n",
      "Round 17, Test Accuracy: 0.9014, Max Acc: 0.9014\n",
      "Round 18, Test Accuracy: 0.9030, Max Acc: 0.9030\n",
      "Round 19, Test Accuracy: 0.9042, Max Acc: 0.9042\n",
      "Round 20, Test Accuracy: 0.9057, Max Acc: 0.9057\n",
      "Round 21, Test Accuracy: 0.9080, Max Acc: 0.9080\n",
      "Round 22, Test Accuracy: 0.9088, Max Acc: 0.9088\n",
      "Round 23, Test Accuracy: 0.9102, Max Acc: 0.9102\n",
      "Round 24, Test Accuracy: 0.9118, Max Acc: 0.9118\n",
      "Round 25, Test Accuracy: 0.9115, Max Acc: 0.9118\n",
      "Round 26, Test Accuracy: 0.9141, Max Acc: 0.9141\n",
      "Round 27, Test Accuracy: 0.9141, Max Acc: 0.9141\n",
      "Round 28, Test Accuracy: 0.9148, Max Acc: 0.9148\n",
      "Round 29, Test Accuracy: 0.9163, Max Acc: 0.9163\n",
      "Round 30, Test Accuracy: 0.9164, Max Acc: 0.9164\n",
      "Round 31, Test Accuracy: 0.9177, Max Acc: 0.9177\n",
      "Round 32, Test Accuracy: 0.9189, Max Acc: 0.9189\n",
      "Round 33, Test Accuracy: 0.9203, Max Acc: 0.9203\n",
      "Round 34, Test Accuracy: 0.9221, Max Acc: 0.9221\n",
      "Round 35, Test Accuracy: 0.9223, Max Acc: 0.9223\n",
      "Round 36, Test Accuracy: 0.9248, Max Acc: 0.9248\n",
      "Round 37, Test Accuracy: 0.9252, Max Acc: 0.9252\n",
      "Round 38, Test Accuracy: 0.9255, Max Acc: 0.9255\n",
      "Round 39, Test Accuracy: 0.9267, Max Acc: 0.9267\n",
      "Round 40, Test Accuracy: 0.9261, Max Acc: 0.9267\n",
      "Round 41, Test Accuracy: 0.9271, Max Acc: 0.9271\n",
      "Round 42, Test Accuracy: 0.9275, Max Acc: 0.9275\n",
      "Round 43, Test Accuracy: 0.9277, Max Acc: 0.9277\n",
      "Round 44, Test Accuracy: 0.9284, Max Acc: 0.9284\n",
      "Round 45, Test Accuracy: 0.9285, Max Acc: 0.9285\n",
      "Round 46, Test Accuracy: 0.9293, Max Acc: 0.9293\n",
      "Round 47, Test Accuracy: 0.9298, Max Acc: 0.9298\n",
      "Round 48, Test Accuracy: 0.9303, Max Acc: 0.9303\n",
      "Round 49, Test Accuracy: 0.9314, Max Acc: 0.9314\n",
      "Round 50, Test Accuracy: 0.9319, Max Acc: 0.9319\n",
      "Round 51, Test Accuracy: 0.9322, Max Acc: 0.9322\n",
      "Round 52, Test Accuracy: 0.9333, Max Acc: 0.9333\n",
      "Round 53, Test Accuracy: 0.9340, Max Acc: 0.9340\n",
      "Round 54, Test Accuracy: 0.9351, Max Acc: 0.9351\n",
      "Round 55, Test Accuracy: 0.9354, Max Acc: 0.9354\n",
      "Round 56, Test Accuracy: 0.9355, Max Acc: 0.9355\n",
      "Round 57, Test Accuracy: 0.9355, Max Acc: 0.9355\n",
      "Round 58, Test Accuracy: 0.9368, Max Acc: 0.9368\n",
      "Round 59, Test Accuracy: 0.9368, Max Acc: 0.9368\n",
      "Round 60, Test Accuracy: 0.9368, Max Acc: 0.9368\n",
      "Round 61, Test Accuracy: 0.9373, Max Acc: 0.9373\n",
      "Round 62, Test Accuracy: 0.9380, Max Acc: 0.9380\n",
      "Round 63, Test Accuracy: 0.9380, Max Acc: 0.9380\n",
      "Round 64, Test Accuracy: 0.9386, Max Acc: 0.9386\n",
      "Round 65, Test Accuracy: 0.9392, Max Acc: 0.9392\n",
      "Round 66, Test Accuracy: 0.9391, Max Acc: 0.9392\n",
      "Round 67, Test Accuracy: 0.9401, Max Acc: 0.9401\n",
      "Round 68, Test Accuracy: 0.9405, Max Acc: 0.9405\n",
      "Round 69, Test Accuracy: 0.9418, Max Acc: 0.9418\n",
      "Round 70, Test Accuracy: 0.9419, Max Acc: 0.9419\n",
      "Round 71, Test Accuracy: 0.9422, Max Acc: 0.9422\n",
      "Round 72, Test Accuracy: 0.9427, Max Acc: 0.9427\n",
      "Round 73, Test Accuracy: 0.9432, Max Acc: 0.9432\n",
      "Round 74, Test Accuracy: 0.9436, Max Acc: 0.9436\n",
      "Round 75, Test Accuracy: 0.9439, Max Acc: 0.9439\n",
      "Round 76, Test Accuracy: 0.9442, Max Acc: 0.9442\n",
      "Round 77, Test Accuracy: 0.9445, Max Acc: 0.9445\n",
      "Round 78, Test Accuracy: 0.9447, Max Acc: 0.9447\n",
      "Round 79, Test Accuracy: 0.9458, Max Acc: 0.9458\n",
      "Round 80, Test Accuracy: 0.9465, Max Acc: 0.9465\n",
      "Round 81, Test Accuracy: 0.9467, Max Acc: 0.9467\n",
      "Round 82, Test Accuracy: 0.9471, Max Acc: 0.9471\n",
      "Round 83, Test Accuracy: 0.9476, Max Acc: 0.9476\n",
      "Round 84, Test Accuracy: 0.9477, Max Acc: 0.9477\n",
      "Round 85, Test Accuracy: 0.9483, Max Acc: 0.9483\n",
      "Round 86, Test Accuracy: 0.9481, Max Acc: 0.9483\n",
      "Round 87, Test Accuracy: 0.9481, Max Acc: 0.9483\n",
      "Round 88, Test Accuracy: 0.9487, Max Acc: 0.9487\n",
      "Round 89, Test Accuracy: 0.9486, Max Acc: 0.9487\n",
      "Round 90, Test Accuracy: 0.9492, Max Acc: 0.9492\n",
      "Round 91, Test Accuracy: 0.9498, Max Acc: 0.9498\n",
      "Round 92, Test Accuracy: 0.9501, Max Acc: 0.9501\n",
      "Round 93, Test Accuracy: 0.9498, Max Acc: 0.9501\n",
      "Round 94, Test Accuracy: 0.9504, Max Acc: 0.9504\n",
      "Round 95, Test Accuracy: 0.9503, Max Acc: 0.9504\n",
      "Round 96, Test Accuracy: 0.9507, Max Acc: 0.9507\n",
      "Round 97, Test Accuracy: 0.9513, Max Acc: 0.9513\n",
      "Round 98, Test Accuracy: 0.9514, Max Acc: 0.9514\n",
      "Round 99, Test Accuracy: 0.9514, Max Acc: 0.9514\n",
      "Round 100, Test Accuracy: 0.9514, Max Acc: 0.9514\n",
      "45.60758261680603\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "begin_time = time.time()\n",
    "round = 1\n",
    "accuracy = []\n",
    "handler.num_clients = trainer.num_clients\n",
    "while handler.if_stop is False:\n",
    "    # server side\n",
    "    sampled_clients = handler.sample_clients()\n",
    "    broadcast = handler.downlink_package\n",
    "\n",
    "    # client side\n",
    "    trainer.local_process(broadcast, sampled_clients)\n",
    "    uploads = trainer.uplink_package\n",
    "\n",
    "    # server side\n",
    "    for pack in uploads:\n",
    "        handler.load(pack)\n",
    "\n",
    "    loss, acc = evaluate(handler._model, nn.CrossEntropyLoss(), test_loader)\n",
    "    accuracy.append(acc)\n",
    "    print(\"Round {}, Test Accuracy: {:.4f}, Max Acc: {:.4f}\".format(round, acc, max(accuracy)))\n",
    "#     if acc>=0.97:\n",
    "#         break\n",
    "    round += 1\n",
    "end_time = time.time()\n",
    "print((end_time - begin_time)/60.0)\n",
    "# torch.save(accuracy, \"./exp_logs/{}, accuracy_{}_B{}_S{}_R{}_Seed{}_T{}.pkl\".format(args.alg, \"mnist\", args.batch_size, args.sample_ratio, args.com_round, args.seed, time.strftime(\"%Y-%m-%d-%H:%M:%S\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "019ae50596e3d4df627f3288be8543f4b17347150bdb9d2aa2e7c637014aee00"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
